feat(save): LazyStore dual-mode WASM backend by dmadisetti · Pull Request #9898 · marimo-team/marimo

dmadisetti · 2026-06-15T23:02:04Z

Summary

This PR creates a WASM compatible LazyStore and introduces the DualMode store concept, where a targeted store can have different behaviors on different platforms.

The LazyLoader wasm dual loads values from an expected cache value using pyodide's native fetch mechanism, and falls back to caching in memory (since no disk writes in WASM)

vercel · 2026-06-15T23:02:10Z

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Jun 30, 2026 7:04pm

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

Flesh out the LazyStore placeholder into a dual-mode store: - LazyStore (native): wraps an inner FileStore, tracks written/touched keys for export. - WasmLazyStore (Pyodide): writes to a shared in-session DictStore; reads fall through to concurrent HTTP fetch from notebook_location()/public/cache/, with path-traversal-safe keys and poisoned-key eviction on corrupt restore. - The single native/WASM decision is made once via a DualLoader registry entry (resolve_loader), so nothing downstream re-checks the platform. Adds WasmExportableStore (export_manifest tracking) + DictStore, the _ACTIVE_LAZY_LOADERS registry, and LazyLoader.flush_all() for export. Also folds in the unserializable-def robustness mechanism: rather than writing a placeholder blob, the loader marks the manifest Item with unserializable_type and reconstructs the UnhashableStub tripwire on load (from_item).

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

Copilot

Pull request overview

Adds a WASM-compatible backend for the lazy persistence layer by introducing a “dual-mode” loader concept (native vs Pyodide/WASM), plus new manifest markers to safely represent unserializable defs without writing placeholder blobs. This fits into marimo._save by extending the existing lazy cache loader/store to support Pyodide constraints (no threads, no durable filesystem) and safer restore semantics.

Changes:

Introduces DualLoader/resolve_loader so PERSISTENT_LOADERS["lazy"] can resolve to LazyLoader (native) or WasmLazyLoader (WASM) via a single is_pyodide() check.
Adds LazyStore/WasmLazyStore and a minimal DictStore, plus a WasmExportableStore interface for batch blob retrieval and export manifest tracking.
Adds Item.unserializable_type + from_item(..., var_name=...) to reconstruct UnhashableStub tripwires from the manifest when serialization fails; expands tests around these behaviors (including torch .pt codec coverage).

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/_save/stubs/test_unhashable_stub.py	Adds tests for explicit `type_name` handling and reconstructing `UnhashableStub` from manifest `Item`.
tests/_save/stubs/test_lazy_codecs.py	New torch-based tests covering lazy codec selection and `.pt` round-trips.
tests/_save/loaders/test_loader.py	Adds regression tests ensuring unserializable values don’t write blobs and restore as tripwires; includes UI stale-blob cleanup behavior.
tests/_save/loaders/test_lazy_wasm.py	New tests for `DictStore`, native vs WASM LazyStore behaviors, key sanitization, batch-path loading, and synchronous WASM writes.
marimo/_save/stubs/lazy_stub.py	Extends `Item` with `unserializable_type`; updates `UnhashableStub` to support explicit `type_name` for manifest reconstruction.
marimo/_save/stores/store.py	Adds `WasmExportableStore` interface with `get_batch()` and `export_manifest()`.
marimo/_save/stores/dict_store.py	New in-memory `DictStore` for WASM sessions.
marimo/_save/stores/init.py	Exports `WasmExportableStore`.
marimo/_save/save.py	Resolves loader registry entries via `resolve_loader(...)` before calling `.partial(...)`.
marimo/_save/loaders/lazy.py	Major update: adds `LazyStore`/`WasmLazyStore`, loader instance registry + `flush_all`, batch blob read path, unserializable markers, and WASM variants for read/write/eviction behavior.
marimo/_save/loaders/init.py	Adds `DualLoader` + `resolve_loader`, registers lazy loader as native/WASM pair.

+        key = self._sanitize_key(key)
+        url = f"{self._base_url()}/{key}"
+        try:
+            with urllib.request.urlopen(url) as resp:
+                return resp.read() if resp.status == 200 else None
+        except Exception:
+            return None


+        try:
+            results = loop.run_until_complete(_fetch_all())
+        except Exception:
+            # run_until_complete on the live pyodide loop requires JSPI
+            # (WebAssembly stack switching), which e.g. Firefox lacks. Fall
+            # back to sequential synchronous XHR via the pyodide_http-patched
+            # urllib — legal in a worker.
+            results = [(k, self._http_get(k)) for k in keys_list]
+        yield from results


…TP reads - Move get_batch (sequential default) and export_keys (empty default, renamed from export_manifest) onto the base Store; drop the WasmExportableStore ABC. LazyStore overrides export_keys to report tracked keys; WasmLazyStore overrides get_batch for concurrent fetch. - WasmLazyStore._http_get/_http_get_batch now cache successful fetches into the inner store and mark them touched, so repeat reads stay local and the export manifest covers HTTP-fetched blobs. - _sanitize_key validates (.. / absolute) but returns the original key instead of a normalized string, avoiding requested/stored key desync. - Remove a duplicate TYPE_CHECKING import block.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

mscolnick · 2026-06-30T13:54:33Z

+
+# Registry of active loaders by name, so a recreated loader reuses the same
+# store (and so `flush_all`/export can reach every session loader).
+_ACTIVE_LAZY_LOADERS: dict[str, LazyLoader] = {}


does this need to be keyed by nested notebooks to avoid conflicts?

mscolnick · 2026-06-30T13:55:12Z

+    def _dispatch_write(self, write_fn: Callable[[], None]) -> None:
+        write_fn()  # synchronous — no threads in Pyodide
+
+    def _on_restore_failure(


missing a test

The corrupt-restore eviction + key-poisoning path was untested. Covers: evicting and poisoning the manifest plus its referenced blobs, the return_value reference, None and undecodable manifests, and an end-to-end corrupt restore through load_cache.

Move the _ACTIVE_LAZY_LOADERS, _POISONED_KEYS, and _WASM_DICT_STORE module globals onto the per-session CacheState via a _cache_state() accessor, falling back to a process-local CacheState when no runtime context is installed (standalone @mo.cache/persistent_cache). The accessor walks to the root context so embedded apps (child kernel contexts) keep sharing the session's loader registry, poison set, and WASM store rather than reinitializing them per child.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

dmadisetti · 2026-06-30T18:32:02Z

@mscolnick moved globals to existing shared object and added test

github-actions · 2026-06-30T19:25:50Z

🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.12-dev26

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

Add CachedLifecycle: an executor lifecycle that skips a cell's body on a cache hit and backfills its defs on a miss. Detection of upstream unserializable defs is done at use-site via a duck-typed `__marimo_unhashable__` marker check (no import coupling to the serialization toolkit), and a pre-flight ref scan requeues the producing cells via a soft MarimoCancelCellError + Scheduler.requeue_for_rerun rather than hard-failing. Stacked on the LazyStore dual-mode backend (#9898): relies on that PR's mark-don't-write mechanism so a cell's own unserializable def restores as an UnhashableStub tripwire instead of raising a PicklingError.

When `marimo export html-wasm --execute` runs, the executed session's LazyLoader writes an export manifest at kernel teardown listing the cache keys it produced. The export step copies exactly those files into `<out_dir>/public/cache/`, where the WASM store's HTTP fallback fetches them — so a cached notebook ships its caches and skips recomputation in the browser. Stacked on the LazyStore dual-mode backend (#9898), which provides `LazyLoader.flush_all()`, the `_ACTIVE_LAZY_LOADERS` registry, and `WasmExportableStore.export_manifest()`.

dmadisetti added the enhancement New feature or request label Jun 15, 2026

vercel Bot deployed to Preview June 15, 2026 23:03 View deployment

vercel Bot deployed to Preview June 15, 2026 23:06 View deployment

vercel Bot deployed to Preview June 15, 2026 23:12 View deployment

dmadisetti force-pushed the dm/save-stubs branch from c2b8ae6 to 7c9f48e Compare June 17, 2026 18:18

Base automatically changed from dm/save-stubs to main June 24, 2026 17:09

dmadisetti force-pushed the dm/lazy-store-merge branch from b4e80dc to 6cb26c5 Compare June 24, 2026 22:50

vercel Bot deployed to Preview June 24, 2026 22:51 View deployment

dmadisetti force-pushed the dm/lazy-store-merge branch from 6cb26c5 to 45fbfde Compare June 24, 2026 23:20

vercel Bot deployed to Preview June 24, 2026 23:21 View deployment

dmadisetti changed the title ~~feat(save): LazyStore dual-mode backend + torch .pt codec~~ feat(save): LazyStore dual-mode WASM backend Jun 25, 2026

dmadisetti requested a review from Copilot June 25, 2026 18:06

Copilot started reviewing on behalf of dmadisetti June 25, 2026 18:06 View session

Copilot AI reviewed Jun 25, 2026

View reviewed changes

dmadisetti requested a review from mscolnick June 25, 2026 18:27

mscolnick reviewed Jun 25, 2026

View reviewed changes

Comment thread marimo/_save/stores/store.py Outdated

mscolnick reviewed Jun 25, 2026

View reviewed changes

Comment thread marimo/_save/loaders/lazy.py Outdated

mscolnick reviewed Jun 25, 2026

View reviewed changes

Comment thread marimo/_save/loaders/lazy.py

vercel Bot deployed to Preview June 29, 2026 17:31 View deployment

dmadisetti marked this pull request as ready for review June 29, 2026 19:03

dmadisetti requested a review from mscolnick June 29, 2026 19:03

mscolnick reviewed Jun 30, 2026

View reviewed changes

Comment thread marimo/_save/loaders/lazy.py Outdated

mscolnick previously approved these changes Jun 30, 2026

View reviewed changes

mscolnick reviewed Jun 30, 2026

View reviewed changes

dmadisetti added 2 commits June 30, 2026 10:06

dmadisetti dismissed mscolnick’s stale review via 76768b2 June 30, 2026 18:26

vercel Bot deployed to Preview June 30, 2026 18:27 View deployment

dmadisetti requested a review from mscolnick June 30, 2026 18:32

mscolnick previously approved these changes Jun 30, 2026

View reviewed changes

comments: cleanup brittle langauge

4065bb6

dmadisetti dismissed mscolnick’s stale review via 4065bb6 June 30, 2026 19:02

ruff: format

6b1e21f

vercel Bot deployed to Preview June 30, 2026 19:04 View deployment

dmadisetti merged commit 9117006 into main Jun 30, 2026
43 checks passed

dmadisetti deleted the dm/lazy-store-merge branch June 30, 2026 19:20

dmadisetti mentioned this pull request Jul 2, 2026

fix(save): keep lazy cache restore robust when a cached dep is absent #10065

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(save): LazyStore dual-mode WASM backend#9898

feat(save): LazyStore dual-mode WASM backend#9898
dmadisetti merged 6 commits into
mainfrom
dm/lazy-store-merge

dmadisetti commented Jun 15, 2026 •

edited

Loading

Uh oh!

vercel Bot commented Jun 15, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mscolnick Jun 30, 2026

Uh oh!

Uh oh!

mscolnick Jun 30, 2026

Uh oh!

dmadisetti commented Jun 30, 2026

Uh oh!

Uh oh!

github-actions Bot commented Jun 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

dmadisetti commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Uh oh!

vercel Bot commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mscolnick Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mscolnick Jun 30, 2026

Choose a reason for hiding this comment

Uh oh!

dmadisetti commented Jun 30, 2026

Uh oh!

Uh oh!

github-actions Bot commented Jun 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dmadisetti commented Jun 15, 2026 •

edited

Loading

vercel Bot commented Jun 15, 2026 •

edited

Loading